import crawler
import requests
import config
import importlib
import reporter
import click
import datetime
import time
import get_info


@click.group()
def control():
    pass


@click.command()
@click.option('--target', '-t', type=str, help='目标地址的url')
@click.option('--is_crawler', '-c', type=bool, default=True, help='是否启用爬虫')
@click.option('--depth', type=int, default=config.default_depth, help='爬取深度设置')
@click.option('--login/--no-login', '-l', default=False, help='是否登录')
@click.option('--report/--no-report', '-r/-nr', default=True, help='是否生成报告')
def scan(target, is_crawler, depth, login, report):
    data = {'beginTime': str(datetime.datetime.now()), 'testName': str(datetime.datetime.now().date()), 'high_vun': 0,
            'mid_vun': 0}
    s_time = time.time()
    print('[+]开始扫描页面，设置为：')
    print('\t-->目标url:' + target)
    print('web服务', get_info.WebServer)
    print("操作系统", get_info.OS)
    print("脚本语言", get_info.Language)

    login_info = None
    if login:
        login_info = config.login_info
        print(f'\t-->登录信息{login_info}')
    if is_crawler:
        print('\t-->启用爬虫', f',爬取深度{depth}')
        c = crawler.Crawler(target, login_info=login_info, depth=depth)
        c.run()
        print(f'[+]爬取完成,共爬取到{len(c.queue)}个网站')

        testResult = []
        for mod in config.scan_module_list:
            ret = importlib.import_module('modules.' + mod)
            print(f'正在执行{mod}扫描')
            for i in c.queue:
                testResult += ret.scan(i, session=c.session)

        # 报告内容设置
        if report:
            print('[-]扫描完成,报告生成中...')
            data['testResult'] = testResult
            data['url_count'] = len(c.queue)
            for t in testResult:
                if t['status'] == '高危':
                    data['high_vun'] += 1
                if t['status'] == '中危':
                    data['mid_vun'] += 1
            data['no_vun'] = data['url_count'] - data['mid_vun'] - data['high_vun']
            data['totalTime'] = time.time() - s_time

            path = reporter.generate_report(data)
            print(f'[+]生成成功,位于{path}')


@click.command()
@click.option('--target', '-t', type=str, help='目标地址的url')
def attack(target):
    print('[+]未发现进攻模块')


control.add_command(scan)
control.add_command(attack)
if __name__ == '__main__':
    control()
